from camel.agents import ChatAgent
from camel.messages import BaseMessage
from camel.models import ModelFactory
from camel.types import ModelPlatformType, RoleType
# 1. 导入依赖库
from py2neo import Graph
from dotenv import load_dotenv
import pymysql
from scripts.regsetup import description

from Task2SQL.Final_Version import TABLE_RELATIONSHIPS, NEO4J_CONFIG

load_dotenv()
import os

MYSQL_CONFIG = {
    "host": "localhost",
    "port": 3306,  # MySQL默认端口，根据实际情况修改
    "user": os.getenv('MYSQL_USER'),
    "password": os.getenv('MYSQL_PASSWORD'),  # 替换为你的MySQL密码
    "database": os.getenv('DATABASE'),  # 数据库名称
    "charset": "utf8mb4",
}


def connect_mysql():
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        print('--->数据库连接成功<---')
        return conn
    except Exception as e:
        raise Exception(f"数据库连接失败: {str(e)}")


'''
提取数据库的表结构信息
'''


def get_table_structure(conn):
    """定义游标，操作数据库"""
    tables = {}  # 用tables存储表的结构信息
    with conn.cursor() as cursor:
        '''with是为了方便管理，在结束后会关闭cursor游标'''
        cursor.execute("SHOW TABLES")
        table_names = [row[0] for row in cursor.fetchall()]
        '''['current_dept_emp', 'departments', 'dept_emp', 'dept_emp_latest_date', 'dept_manager', 'employees', 'salaries', 'titles']'''

        if not table_names:
            raise Exception("--->未获取到表名，请确认已经建表<---")

        '''提取字段信息'''
        for table_name in table_names:
            cursor.execute(f'show columns from {table_name}')
            columns = cursor.fetchall()
            '''
            (('emp_no', 'int(11)', 'NO', '', None, ''), ('dept_no', 'char(4)', 'NO', '', None, ''), ('from_date', 'date', 'YES', '', None, ''), ('to_date', 'date', 'YES', '', None, ''))
            (('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('dept_name', 'varchar(40)', 'NO', 'UNI', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', '', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', '', None, ''), ('from_date', 'date', 'YES', '', None, ''), ('to_date', 'date', 'YES', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('dept_no', 'char(4)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', '', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('birth_date', 'date', 'NO', '', None, ''), ('first_name', 'varchar(14)', 'NO', '', None, ''), ('last_name', 'varchar(16)', 'NO', '', None, ''), ('gender', "enum('M','F')", 'NO', '', None, ''), ('hire_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('salary', 'int(11)', 'NO', '', None, ''), ('from_date', 'date', 'NO', 'PRI', None, ''), ('to_date', 'date', 'NO', '', None, ''))
            (('emp_no', 'int(11)', 'NO', 'PRI', None, ''), ('title', 'varchar(50)', 'NO', 'PRI', None, ''), ('from_date', 'date', 'NO', 'PRI', None, ''), ('to_date', 'date', 'YES', '', None, ''))
             '''

            fields = []
            for col in columns:
                '''
                field_name用来获取每个字段的字段名
                key_type用来判断该字段是否为主键或者外键
                '''
                field_name = col[0]
                key_type = col[3]
                if key_type == 'PRI':
                    fields.append(f"{field_name} [主键]")
                elif key_type == 'MUL':
                    fields.append(f"{field_name} [外键]")
                else:
                    fields.append(field_name)
                    '''
                    ['emp_no', 'dept_no', 'from_date', 'to_date']
                    ['dept_no [主键]', 'dept_name']
                    ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']
                    ['emp_no', 'from_date', 'to_date']
                    ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']
                    ['emp_no [主键]', 'birth_date', 'first_name', 'last_name', 'gender', 'hire_date']
                    ['emp_no [主键]', 'salary', 'from_date [主键]', 'to_date']
                    ['emp_no [主键]', 'title [主键]', 'from_date [主键]', 'to_date']
                    '''
            tables[table_name] = {
                'name': table_name,
                'fields': fields,
            }
            '''
            {'name': 'current_dept_emp', 'fields': ['emp_no', 'dept_no', 'from_date', 'to_date']}
            {'name': 'departments', 'fields': ['dept_no [主键]', 'dept_name']}
            {'name': 'dept_emp', 'fields': ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']}
            {'name': 'dept_emp_latest_date', 'fields': ['emp_no', 'from_date', 'to_date']}
            {'name': 'dept_manager', 'fields': ['emp_no [主键]', 'dept_no [主键]', 'from_date', 'to_date']}
            {'name': 'employees', 'fields': ['emp_no [主键]', 'birth_date', 'first_name', 'last_name', 'gender', 'hire_date']}
            {'name': 'salaries', 'fields': ['emp_no [主键]', 'salary', 'from_date [主键]', 'to_date']}
            {'name': 'titles', 'fields': ['emp_no [主键]', 'title [主键]', 'from_date [主键]', 'to_date']}
            '''
    print(f"--->成功提取到{len(tables)}张表的结构信息<---")
    return tables


'''
表间关系定义--->人工补充业务语义
'''
TABLE_RELATIONSHIPS = [
    {
        "from_table": "employees",
        "to_table": "dept_emp",
        "description": "员工隶属于部门（多对多关联，通过中间表dept_emp实现）",
        "field_relation": "emp_no → dept_emp.emp_no"
    },
    {
        "from_table": "departments",
        "to_table": "dept_emp",
        "description": "部门包含员工（多对多关联，通过中间表dept_emp实现）",
        "field_relation": "dept_no → dept_emp.dept_no"
    },
    {
        "from_table": "employees",
        "to_table": "dept_manager",
        "description": "员工担任部门经理（多对多关联，通过中间表dept_manager实现）",
        "field_relation": "emp_no → dept_manager.emp_no"
    },
    {
        "from_table": "departments",
        "to_table": "dept_manager",
        "description": "部门拥有经理（多对多关联，通过中间表dept_manager实现）",
        "field_relation": "dept_no → dept_manager.dept_no"
    },
    {
        "from_table": "employees",
        "to_table": "salaries",
        "description": "员工拥有薪资记录（一对多关联，一个员工可有多条薪资记录）",
        "field_relation": "emp_no → salaries.emp_no"
    },
    {
        "from_table": "employees",
        "to_table": "titles",
        "description": "员工拥有职位头衔记录（一对多关联，一个员工可有多条职位记录）",
        "field_relation": "emp_no → titles.emp_no"
    },
    {
        "from_table": "dept_emp",
        "to_table": "dept_emp_latest_date",
        "description": "部门员工关联表数据生成员工最新部门信息视图（视图依赖基础表）",
        "field_relation": "emp_no,from_date,to_date → dept_emp_latest_date.emp_no,from_date,to_date"
    },
    {
        "from_table": "dept_emp",
        "to_table": "current_dept_emp",
        "description": "部门员工关联表数据参与生成当前部门员工视图（视图依赖基础表）",
        "field_relation": "emp_no,dept_no,from_date,to_date → current_dept_emp.emp_no,dept_no,from_date,to_date"
    },
    {
        "from_table": "dept_emp_latest_date",
        "to_table": "current_dept_emp",
        "description": "员工最新部门信息视图数据参与生成当前部门员工视图（视图依赖视图）",
        "field_relation": "emp_no,from_date,to_date → current_dept_emp.emp_no,from_date,to_date"
    }
]

NEO4J_CONFIG ={
    'uri': 'bolt://localhost:7687',
    'user': 'neo4j',
    'password': os.getenv('PASSWORD_NEO')
}
'''
操作Neo4j数据库
'''


def connect_neo4j():
    try:
        graph = Graph(
            uri=NEO4J_CONFIG['uri'],
            auth=(NEO4J_CONFIG['user'], NEO4J_CONFIG['password'])
        )
        print('--->连接Neo4j成功<---')
        return graph
    except Exception as e:
        raise Exception(f"Neo4j连接失败:{e}")


'''
创建唯一性约束
'''


def create_neo4j_constraints(graph):
    graph.run("CREATE CONSTRAINT IF NOT EXISTS FOR (t:Table) REQUIRE t.name IS UNIQUE")
    print("--->唯一性约束创建成功<---")


'''
创建节点
'''


def create_table_nodes(graph, tables):
    for table in tables.values():
        graph.run(
            """
            merge (t:Table {name:$name}) set t.label = $label, t.fields = $fields
            """,
            name=table["name"],
            label=table["name"],
            fields=table["fields"],
        )
    print(f"--->成功创建{len(tables)}个节点<---")


'''
创建关系边
'''
def create_table_relationships(graph):
    for rel in TABLE_RELATIONSHIPS:
        graph.run(
            '''
            match (from:Table {name:$from_table})
            match (to:Table {name:$to_table}) 
            merge (from)-[r:REFERENCES {
            description:$description,
            field_relation:$field_relation
            }]->(to)
            ''',
            from_table=rel["from_table"],
            to_table=rel["to_table"],
            field_relation=rel["field_relation"],
            description=rel["description"],
        )
    print(f"--->已在Neo4j中创建 {len(TABLE_RELATIONSHIPS)} 条表关系<---")

'''
查询方法，提供给智能体
'''
def query_graph_relationships(graph, table_names=None):
    """查询Neo4j中表之间的关系,
    这个函数的作用是从Neo4j图数据库中查询指定表之间的关系信息，为Text2SQL系统提供表间关联的语义信息。
    """
    if not table_names:   #如果未指定要查询的表名，则使用默认表名列表
        table_names = ["departments", "dept_emp", "dept_manager", "employees","salaries","titles"]

    query = """
    MATCH (t1:Table)-[r:REFERENCES]-(t2:Table)
    WHERE t1.name IN $table_names AND t2.name IN $table_names AND t1.name < t2.name
    RETURN t1.name AS from_table, r.field_relation AS relationship, t2.name AS to_table
    """

    result = graph.run(query, table_names=table_names).data()
    print('--->成功获取到Neo4j中表之间的关系<---')
    return result

SYSTEM_PROMPT ="""
你是Text2SQL专家，必须严格按以下规则将用户问题转为MySQL语句：
一、工具调用强制格式（必须严格遵守，否则无法执行）
1. 查询所有表结构：
    [get_table_structure_tool()]
2.查询所有表关系：
    [query_graph_relationships_tool()]
4.查询指定表关系(如薪资表和员工表):
    [query_graph_relationships_tool(graph,table_names=["employees","salaries"])
    
二、工作流程
1.第一步，get_table_structure_tool()获取所有表结构（必须执行）
2.第二步，query_graph_relationships_tool()，根据用户的问题，查询指定表之间的关系
3.第三步，基于工具返回的**实际表名和字段**生成SQL，禁止使用默认表名
三、SQL生成要求
- 表名和字段名必须与工具返回的完全一致（例如工具返回"employees"，不可简写）；
- 生成SQL后必须附带说明，解释使用的表、字段和计算逻辑。
"""
def begin(user_question):
    conn = None
    try:
        conn = connect_mysql()
        graph = connect_neo4j()
        table_structure = get_table_structure(conn)
        create_neo4j_constraints(graph)
        create_table_nodes(graph, table_structure)
        create_table_relationships(graph)
        print("\n🎉 Text2SQL语义图谱构建完成！")
        print("   可通过 http://localhost:7474 访问Neo4j浏览器查看图谱")
        print("   在Neo4j浏览器中执行 MATCH (n) RETURN n 可查看所有节点和关系")
        Model_Type = os.getenv("Model_Type")
        Model_Api = os.getenv("ZHIPU_API_KEY")
        Model_Url = os.getenv("MODEL_URL")
        # 修复工具调用时MySQL连接对象传递问题，确保get_table_structure能正确获取连接
        model = ModelFactory.create(
            model_platform=ModelPlatformType.OPENAI,
            model_type=Model_Type,
            api_key=Model_Api,
            url=Model_Url
        )

        # 定义包装器，确保get_table_structure始终使用已建立的mysql_conn
        def get_table_structure_tool():
            """
            工具描述：获取所有表的结构信息，包括表名、字段名、字段类型、主键信息等。
            返回值：以字典形式返回所有表的结构详情，适用于数据库结构理解与SQL生成。
            """
            return get_table_structure(conn)

        # 同理，包装query_graph_relationships，确保传递neo4j_graph
        def query_graph_relationships_tool(table_names=None):
            """
            工具描述：查询指定表之间的关系（如外键、业务关联等），如未指定表名则返回所有表关系。
            参数：table_names（可选）- 需要查询关系的表名列表。
            返回值：返回表与表之间的关系详情，便于理解数据表之间的业务联系。
            """
            return query_graph_relationships(graph, table_names)

        agent = ChatAgent(
            model=model,
            system_message=SYSTEM_PROMPT,
            output_language="中文",
            tools=[query_graph_relationships_tool, get_table_structure_tool]
        )
        msg = BaseMessage(
            role_name="用户",
            role_type=RoleType.USER,
            content=user_question,
            meta_dict={}
        )
        response = agent.step(msg)
        print(f"大模型响应：{response.msg.content}")

    except Exception as e:
        print(f"\n❌ 构建过程出错: {str(e)}")
        raise  # 抛出异常以便调试

    finally:
    # 确保MySQL连接被关闭
        if conn:
            conn.close()
            print("\n🔌 MySQL连接已关闭")

if __name__ == '__main__':
    question = ("查询部门号为d003中，生日为同一天的女性员工的平均薪资"
                "结果显示为：部门名称、性别、生日、平均薪资")
    begin(question)
